import numpy as np
import matplotlib.pyplot as plt

def numerical_diff(f,x):
    h = 1e-4
    return (f(x+h)-f(x-h))/(2*h)

def function_2(x): # x为numpy数组
    return np.sum(x**2)

#1 求x0=3.0,x1=4.0 处，函数关于x0的偏导数，首先把二元函数转换为一元函数
def function_tem1(x0):
    return x0*x0 + 4.0**2.0

# 再使用一元函数数值微分求解
# print(numerical_diff(function_tem1,3.0))

# 2 求x0=3.0,x1=4.0处，函数关于x1的偏导数，也是首先要把二元函数转换为一元函数
def function_tem2(x1):
    return 3.0**2+x1*x1

# 再使用一元函数数值微分求解
# print(numerical_diff(function_tem2,4))

# 梯度计算
def _numerical_gradient_no_batch(f, x):
    h = 1e-4 # 0.0001
    grad = np.zeros_like(x)
    
    for idx in range(x.size):
        tmp_val = x[idx]
        x[idx] = tmp_val + h
        fxh1 = f(x) # f(x+h)
        
        x[idx] = tmp_val - h 
        fxh2 = f(x) # f(x-h)
        grad[idx] = (fxh1 - fxh2) / (2*h)
        
        x[idx] = tmp_val # 还原值
        
    return grad
def numerical_gradient(f,X):
    if X.ndim == 1:
        return _numerical_gradient_no_batch(f,X)
    else:
        grad = np.zeros_like(X)

        for idx, x in enumerate(X):
            grad[idx] = _numerical_gradient_no_batch(f,x)
        return grad
if __name__ == '__main__':
    print(numerical_gradient(function_2,np.array([[3.0,4.0],[0.0,2.0]])))
